#ifndef READMESH_HPP_
#define READMESH_HPP_

#include <filesystem>
#include <fstream>
#include <stdexcept>
#include <string>
#include <vector>

#include <Eigen/Core>


enum class AttribType { POINT, CELL };

enum class AttribFormat { SCALAR, VECTOR };

enum class AttribDataType { INT, DOUBLE };

struct MeshAttrib {
  AttribType type;
  AttribFormat fmt;
  AttribDataType dataType;

  std::string name;
  const void *data;

  const char *getDataTypeName() const {
    switch (dataType) {
    case AttribDataType::INT:
      return "int";
    case AttribDataType::DOUBLE:
      return "double";
    }
    return nullptr;
  }

  void writeScalarDataArray(std::ofstream &fout, int size) const {
    fout << "SCALARS " << name << ' ' << getDataTypeName() << '\n';
    fout << "LOOKUP_TABLE default\n";
    if (dataType == AttribDataType::INT) {
      for (int i = 0; i < size; ++i)
        fout << static_cast<const int *>(data)[i] << '\n';
    } else if (dataType == AttribDataType::DOUBLE) {
      for (int i = 0; i < size; ++i)
        fout << static_cast<const double *>(data)[i] << '\n';
    }
  }

  void writeVectorDataArray(std::ofstream &fout, int size) const {
    fout << "VECTORS " << name << ' ' << getDataTypeName() << '\n';
    if (dataType == AttribDataType::INT) {
      for (int i = 0; i < size; ++i) {
        const int *d = static_cast<const int *>(data) + 3 * i;
        fout << d[0] << ' ' << d[1] << ' ' << d[2] << '\n';
      }
    } else if (dataType == AttribDataType::DOUBLE) {
      for (int i = 0; i < size; ++i) {
        const double *d = static_cast<const double *>(data) + 3 * i;
        fout << d[0] << ' ' << d[1] << ' ' << d[2] << '\n';
      }
    }
  }
};

template <typename DerivedV, typename DerivedT>
inline void writeHexMeshToVtk(const std::filesystem::path &path,
                              const Eigen::MatrixBase<DerivedV> &V,
                              const Eigen::MatrixBase<DerivedT> &T,
                              const std::vector<MeshAttrib> &attribs = {}) {
  std::ofstream fout(path);

  if (!fout.is_open()) {
    throw std::runtime_error("Failed to open file: " + path.string());
  }

  fout << "# vtk DataFile Version 3.0\n";
  fout << "tetmesh\n";
  fout << "ASCII\n";
  fout << "DATASET UNSTRUCTURED_GRID\n";
  fout << "POINTS " << V.rows() << " double\n";
  for (int i = 0; i < V.rows(); ++i) {
    fout << V(i, 0) << " " << V(i, 1) << " " << V(i, 2) << "\n";
  }
  fout << "CELLS " << T.rows() << " " << 9 * T.rows() << "\n";
  for (int i = 0; i < T.rows(); ++i) {
    fout << "8 ";
    for (int j = 0; j < 8; ++j)
      fout << T(i, j) << " ";
    fout << "\n";
  }

  fout << "CELL_TYPES " << T.rows() << "\n";
  for (int i = 0; i < T.rows(); ++i) {
    fout << "12\n";
  }

  bool hasPointAttrib = false;
  bool hasCellAttrib = false;
  for (const auto &a : attribs) {
    if (a.type == AttribType::POINT)
      hasPointAttrib = true;
    else if (a.type == AttribType::CELL)
      hasCellAttrib = true;
  }

  if (hasPointAttrib) {
    fout << "POINT_DATA " << V.rows() << "\n";
    for (const auto &a : attribs) {
      if (a.type != AttribType::POINT)
        continue;
      if (a.fmt == AttribFormat::SCALAR) {
        a.writeScalarDataArray(fout, V.rows());
      } else if (a.fmt == AttribFormat::VECTOR) {
        a.writeVectorDataArray(fout, V.rows());
      }
    }
  }

  if (hasCellAttrib) {
    fout << "CELL_DATA " << T.rows() << "\n";
    for (const auto &a : attribs) {
      if (a.type != AttribType::CELL)
        continue;
      if (a.fmt == AttribFormat::SCALAR) {
        a.writeScalarDataArray(fout, T.rows());
      } else if (a.fmt == AttribFormat::VECTOR) {
        a.writeVectorDataArray(fout, T.rows());
      }
    }
  }
}

#endif